#!/usr/bin/python2.6
from math import sqrt,pow
import xml.parsers.expat
from sqlite3 import dbapi2 as sqlite
import os
import codecs
import sys

"""specify I\O file names here"""
in_file_name = "source.osm"
node_db_file = "osm.db"
#MVC = 500000 #MAX VERTEX COUNT PER BUCKET
MVC = 50000

class Node:
    def __init__(self,id,lat=0,lon=0):
        self.id = id
        self.lat = lat
        self.lon = lon
    lat = 0
    lon = 0
    id = 0L
    def get_edge_tuple(self,second):
        return (self.id, second.id)

class MyXML:
    """class for wrapping expat xml parser"""
    Parser = ""
    xml_file_name = ""
    xml_file = ""
    ValueDict = {}
    # prepare for parsing
    def __init__(self, xml_file_name):
            assert(xml_file_name != "")
            self.xml_file_name = xml_file_name
            self.Parser = xml.parsers.expat.ParserCreate()
            self.Parser.StartElementHandler = self.handleStartElement
            self.Parser.EndElementHandler = self.handleEndElement
#                self.Parser.CharacterDataHandler = self.char_data
#        def char_data(self,data):
#            print 'Character data:', repr(data)
    # parse the XML file
    def Parse(self):
        print "Processing ",self.xml_file_name, " ..."
        self.xml_file = open(self.xml_file_name, "r")
        self.Parser.ParseFile(self.xml_file)
        print "Processing completed"
    def Close(self):
        self.xml_file.close()
        print "File ",self.xml_file_name, " closed"
     # will be overwritten w/ implementation specific methods
    def handleStartElement(self, name, attrs): pass
    def handleEndElement(self, name): pass
    def setStartElementHandler(self,handler):
        MyXML.handleStartElement = handler
        self.Parser.StartElementHandler = self.handleStartElement
    def setEndElementHandler(self,handler):
        MyXML.handleEndElement = handler
        self.Parser.EndElementHandler = self.handleEndElement

def CreateDB(connection):
    cursor = connection.cursor()
    #firstly id is osm_id and after processing id is id in bucket
    cursor.execute("""CREATE TABLE vertices (bucket_id INT, id INT, osm_id BIGINT,lat REAL, lon REAL)""")
    
    #cursor.execute("""CREATE TABLE edges (source_bucket_id INT, target_bucket_id INT, source_id INT, target_id INT, 
    #		      source_osm_id BIGINT, target_osm_id BIGINT, highway_tag VARCHAR,
    #                                     PRIMARY KEY(source_osm_id,target_osm_id))""")
    cursor.execute("""CREATE TABLE edges (source_osm_id BIGINT, target_osm_id BIGINT, highway_tag VARCHAR,
                                         PRIMARY KEY(source_osm_id,target_osm_id))""")        
    cursor.execute("""CREATE TABLE buckets (id INT PRIMARY KEY, min_lat REAL, min_lon REAL, max_lat REAL, max_lon REAL, 
					  vertex_count INT, edge_count INT, info BLOB, vertices BLOB, edges BLOB)""")

    #cursor.execute("""CREATE INDEX ix_edges1 ON edges(source_id)""")
    #cursor.execute("""CREATE INDEX ix_edges2 ON edges(target_id)""")
    cursor.execute("""CREATE INDEX ix_vertices_id ON vertices(osm_id)""")
    
    connection.commit()

def Handler_SNd(self,name, attrs):
        if name == "node":
	    lat = float(attrs["lat"])
	    #if lat < 80 and lat > -80:
	    lon = float(attrs["lon"])
	    connection = self.ValueDict["connection"]
	    cursor = self.ValueDict["cursor"]
	    ins_params = (long(attrs["id"]), lat,lon)
	    cursor.execute("""INSERT OR IGNORE INTO vertices (osm_id,lat,lon) VALUES (?, ?, ?)""", ins_params)

def Handler_SWay(self,name, attrs):
    if self.ValueDict["in_way"] and name == "nd":
        self.ValueDict["way_nodes"].append(long(attrs["ref"]))

    elif name == "way":
        self.ValueDict["in_way"] = 1

    elif self.ValueDict["in_way"] and name == "tag":
           if attrs["k"] == "highway":
               self.ValueDict["highway"] = unicode(attrs["v"])
           if attrs["k"] == "oneway":
               self.ValueDict["oneway"] = unicode(attrs["v"])

def Handler_EWay(self,name):
    """When we parsed the way,let's write resulting edges to db"""
    if self.ValueDict["in_way"] and name == "way":
        way_nodes_ids = self.ValueDict["way_nodes"]
        if self.ValueDict["highway"]:
            cursor = Trans.ValueDict["cursor"]
            for i in xrange(len(way_nodes_ids)-1):                
                #self.ValueDict["e_count"] += 1
		if ("oneway" not in self.ValueDict): #ok, then two edges for two directions
		    edge_tuple = (way_nodes_ids[i] , way_nodes_ids[i+1], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)
		    edge_tuple = (way_nodes_ids[i+1], way_nodes_ids[i], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)			
		elif (self.ValueDict["oneway"] == "-1"): #reverse edges
                    edge_tuple = (way_nodes_ids[i+1], way_nodes_ids[i], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)
		elif (self.ValueDict["oneway"] == "yes"): #single direction
                    edge_tuple = (way_nodes_ids[i] , way_nodes_ids[i+1], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)				
		else: #two edges for two directions
                    edge_tuple = (way_nodes_ids[i] , way_nodes_ids[i+1], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)
		    edge_tuple = (way_nodes_ids[i+1], way_nodes_ids[i], self.ValueDict["highway"])
		    cursor.execute("""INSERT OR IGNORE INTO edges (source_osm_id, target_osm_id, highway_tag ) VALUES (?, ?, ?)""", edge_tuple)		
            
        self.ValueDict["in_way"] = 0
        self.ValueDict["highway"] = None
	self.ValueDict["oneway"] = "no"
        del way_nodes_ids[:]
        way_nodes_ids[:] = []
def DeleteUnNeededVerticesAndEdges(cursor):
   """There is a heuristic: each edge references to target vertex, but next edge references this vertex in source node
      so, we can delete only by source_id - BUT NO IT DOESN'T WORK. Example: visiachaia vershina""" 
   #delete edges and vertice with high latitude
   cursor.execute("BEGIN TRANSACTION") 
   cursor.execute("""DELETE FROM edges 
		     WHERE source_osm_id IN (SELECT osm_id FROM vertices WHERE lat > 80 OR lat < -80)""")
   cursor.execute("""DELETE FROM vertices WHERE lat > 80 OR lat < -80""")
   #delete unused vertices
   cursor.execute("""DELETE FROM vertices WHERE osm_id NOT IN (SELECT source_osm_id FROM edges UNION  SELECT target_osm_id FROM edges) """)   
   #cursor.execute("COMMIT TRANSACTION")
   connection.commit()


def GetLatLonSize(max_lat, min_lat, max_lon, min_lon, current_bucket_max_count):
   dim1_bucket_count = int(sqrt(current_bucket_max_count)) + 1
   return ( (max_lat - min_lat)/dim1_bucket_count, (max_lon - min_lon)/dim1_bucket_count )

def MakeBuckets(cursor):
   cursor.execute("""SELECT count(*) FROM vertices""")
   row = cursor.fetchone() 
   #print row
   vertex_count = long(row[0])
   cursor.execute("""SELECT min(lat) FROM vertices""")
   row = cursor.fetchone()
   #print row
   min_lat = float(row[0])
   cursor.execute("""SELECT min(lon) FROM vertices""")
   row = cursor.fetchone()
   min_lon = float(row[0])
   cursor.execute("""SELECT max(lat) FROM vertices""")
   row = cursor.fetchone()
   max_lat = float(row[0])
   cursor.execute("""SELECT max(lon) FROM vertices""")
   row = cursor.fetchone()
   max_lon = float(row[0])
   #print "min_lat ", min_lat, "min_lon: ", min_lon, "max_lat: ", max_lat, "max_lon: ", max_lon, "vertex_count: ", vertex_count
   #cursor.execute("BEGIN TRANSACTION")
   current_bucket_max_count = int(vertex_count/MVC) + 1
   current_MVC = MVC + 1
   current_bucket_lat_size = GetLatLonSize(max_lat, min_lat, max_lon, min_lon, current_bucket_max_count)[0]
   current_bucket_lon_size = GetLatLonSize(max_lat, min_lat, max_lon, min_lon, current_bucket_max_count)[1]
   lat_max_number = long((max_lat - min_lat)/current_bucket_lat_size) + 1
   lon_max_number = long((max_lon - min_lon)/current_bucket_lon_size) + 1

   while(current_MVC > MVC):
       cursor.execute("SELECT max(cnt) FROM " +
                      "(SELECT CAST( (lat - ("+str(min_lat)+"))/" + str(current_bucket_lat_size) + " AS BIGINT ) * (" + str(lon_max_number) +
                      ") + CAST( (lon - ("+str(min_lon)+"))/(" + str(current_bucket_lon_size) + ") AS BIGINT ) AS temp_bucket_id ,count(*) as cnt FROM vertices " +
                      "GROUP BY temp_bucket_id)")
       row = cursor.fetchone()
       current_MVC = long(row[0])

       current_bucket_max_count = int(current_bucket_max_count * (float(current_MVC)/MVC)*1.1) + 1
       current_bucket_lat_size = GetLatLonSize(max_lat, min_lat, max_lon, min_lon, current_bucket_max_count)[0]
       current_bucket_lon_size = GetLatLonSize(max_lat, min_lat, max_lon, min_lon, current_bucket_max_count)[1]
       lat_max_number = long((max_lat - min_lat)/current_bucket_lat_size) + 1
       lon_max_number = long((max_lon - min_lon)/current_bucket_lon_size) + 1
       print "current max vert count: ", current_MVC, "MVC0: ", MVC, "current predicted bucket max count: ", current_bucket_max_count
       #print "lat_max_number", lat_max_number, current_bucket_lat_size, current_bucket_lon_size
   
   cursor.execute("UPDATE vertices set bucket_id = CAST( (lat - ("+str(min_lat)+"))/(" + str(current_bucket_lat_size) + ") AS BIGINT ) * (" + str(lon_max_number) +
                  ") + CAST( (lon - ("+str(min_lon)+"))/(" + str(current_bucket_lon_size) + ") AS BIGINT )")
   
   for lat_n in xrange(0, lat_max_number):
       for lon_n in xrange(0, lon_max_number):
           cursor.execute("""INSERT INTO buckets (id, min_lat, min_lon, max_lat, max_lon) 
                             VALUES(?,?,?,?,?)""", 
                             ( lat_n*lon_max_number + lon_n, lat_n*current_bucket_lat_size + min_lat, \
                             lon_n*current_bucket_lon_size + min_lon, (lat_n+1)*current_bucket_lat_size + min_lat, (lon_n+1)*current_bucket_lon_size + min_lon) 
                         )
           #print ( lat_n*lon_max_number + lon_n, lat_n*current_bucket_lat_size + min_lat, \
           #                  lon_n*current_bucket_lon_size + min_lon, (lat_n+1)*current_bucket_lat_size + min_lat, (lon_n+1)*current_bucket_lon_size + min_lon)
   #cursor.execute("COMMIT TRANSACTION")
   cursor.execute("""CREATE INDEX ix_vertices_bucket_id ON vertices(bucket_id)""")
   cursor.connection.commit() 

def SetNewIds(cursor):
    cursor_upd = cursor.connection.cursor()
    cursor.execute("SELECT id from buckets")
    temp = cursor.fetchall()
    buckets = {}
    for id in temp:
        buckets[long(id[0])] = 0
    #print "Bucket count: ", len(buckets)
    #print "MVC ", MVC

    cursor_upd.execute("BEGIN TRANSACTION")
    cursor.execute("SELECT bucket_id, osm_id FROM vertices")
	
    vertex = cursor.fetchone()  
    while vertex:           
	#print "vertex:", vertex       
        cursor_upd.execute("UPDATE vertices set id = :new_id where osm_id = :old_id",
                               {'new_id': buckets[vertex[0]], 'old_id':vertex[1]})
	buckets[vertex[0]] += 1
        vertex = cursor.fetchone()        

    #print "vertex:", vertex
    for buck_id, count in buckets.iteritems():
 	cursor_upd.execute("UPDATE buckets set vertex_count = ? where id = ?", (count,buck_id))
    
    #cursor.execute("COMMIT TRANSACTION")
    cursor_upd.connection.commit() 
    #for item in buckets:
    #     print item

#checking console args
if sys.argv[1]:
    in_file_name = sys.argv[1]
if sys.argv[2]:
    node_db_file = sys.argv[2]

#creating new DB file
if os.path.exists(node_db_file):
    os.remove(node_db_file)
    print "old db_file removed"

print "creating db..."
connection = sqlite.connect(node_db_file) #, isolation_level = None
cursor = connection.cursor()
CreateDB(connection)
print "creating db done..."

#MakeBuckets(cursor)

print "inserting all nodes to sql lite db"
Trans = MyXML(in_file_name)
Trans.setStartElementHandler(Handler_SNd)
Trans.setEndElementHandler(None)
Trans.ValueDict["connection"] = connection
Trans.ValueDict["cursor"] = cursor
cursor.execute("BEGIN TRANSACTION")
Trans.Parse()
#cursor.execute("COMMIT TRANSACTION")
Trans.Close()
connection.commit()
print "inserting nodes done..."

#cursor.execute("SELECT * FROM vertices LIMIT 10")
#print cursor.fetchall()

print "inserting edges..."
Trans = None
Trans = MyXML(in_file_name)
Trans.ValueDict["connection"] = connection
Trans.ValueDict["cursor"] = cursor
Trans.ValueDict["in_way"] = None
Trans.ValueDict["highway"] = None
Trans.ValueDict["way_nodes"] = []
Trans.ValueDict["e_count"] = 0
Trans.setStartElementHandler(Handler_SWay)
Trans.setEndElementHandler(Handler_EWay)
cursor.execute("BEGIN TRANSACTION")
Trans.Parse()
#cursor.execute("COMMIT TRANSACTION")
Trans.Close()
connection.commit()
print "inserting edges done..."

#cursor.execute("SELECT * FROM edges LIMIT 10")
#print cursor.fetchall()

print "deleting edges and vertices with high latitude..."
DeleteUnNeededVerticesAndEdges(cursor)
print "deleting edges and vertices with high latitude done..."

#cursor.execute("SELECT * FROM vertices LIMIT 10")
#print cursor.fetchall()
#cursor.execute("SELECT * FROM edges LIMIT 10")
#print cursor.fetchall()

print "making buckets..."
MakeBuckets(cursor)
print "making buckets done..."

print "setting new ids..."
SetNewIds(cursor)
print "setting new ids done..."

#print "Vertices:"
#cursor.execute("SELECT * FROM vertices LIMIT 10")
#print cursor.fetchall()
#print "Buckets:"
#cursor.execute("SELECT * FROM buckets ORDER BY id LIMIT 30")
#print cursor.fetchall()
#print "Bucket ids from vertices:"
#cursor.execute("SELECT bucket_id, count(*), min(lat), max(lat), min(lon), max(lon) FROM vertices GROUP BY bucket_id")
#print cursor.fetchall()

print "Max count in buckets:"
cursor.execute("SELECT max(cnt) FROM " +
                      "(SELECT bucket_id, count(*) as cnt FROM vertices " +
                      "GROUP BY bucket_id)")
print cursor.fetchall()
print "Bucket count:"
cursor.execute("SELECT count(*) FROM buckets")
print cursor.fetchall()

cursor.close()
connection.close()
print "ALL work is done"
